#!/usr/bin/env python3
"""
CNN图像识别器测试
"""

import sys
import os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

import pytest
import torch
import numpy as np
from PIL import Image

from models.image_classifier import CNNImageClassifier, ImageProcessor, ImageRecognizer


class TestCNNImageClassifier:
    """测试CNN图像分类器"""
    
    def test_model_initialization(self):
        """测试模型初始化"""
        model = CNNImageClassifier(num_classes=10)
        
        assert model is not None
        assert isinstance(model, torch.nn.Module)
        assert hasattr(model, 'conv1')
        assert hasattr(model, 'fc2')
    
    def test_forward_pass(self):
        """测试前向传播"""
        model = CNNImageClassifier(num_classes=10)
        
        # 创建测试输入
        batch_size = 2
        input_tensor = torch.randn(batch_size, 3, 32, 32)
        
        # 前向传播
        output = model(input_tensor)
        
        # 检查输出形状
        assert output.shape == (batch_size, 10)
        assert not torch.isnan(output).any()
        assert not torch.isinf(output).any()
    
    def test_different_num_classes(self):
        """测试不同类别数量的模型"""
        for num_classes in [2, 10, 100]:
            model = CNNImageClassifier(num_classes=num_classes)
            
            input_tensor = torch.randn(1, 3, 32, 32)
            output = model(input_tensor)
            
            assert output.shape == (1, num_classes)


class TestImageProcessor:
    """测试图像处理器"""
    
    def test_processor_initialization(self):
        """测试处理器初始化"""
        processor = ImageProcessor(image_size=(32, 32))
        
        assert processor is not None
        assert hasattr(processor, 'transform')
    
    def test_process_numpy_array(self):
        """测试处理numpy数组"""
        processor = ImageProcessor(image_size=(32, 32))
        
        # RGB数组
        rgb_array = np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8)
        processed = processor.process(rgb_array)
        
        assert processed.shape == (1, 3, 32, 32)
        assert processed.dtype == torch.float32
    
    def test_process_grayscale_array(self):
        """测试处理灰度数组"""
        processor = ImageProcessor(image_size=(32, 32))
        
        # 灰度数组
        gray_array = np.random.randint(0, 255, (64, 64), dtype=np.uint8)
        processed = processor.process(gray_array)
        
        assert processed.shape == (1, 3, 32, 32)  # 应该转换为3通道
    
    def test_process_tensor(self):
        """测试处理PyTorch tensor"""
        processor = ImageProcessor(image_size=(32, 32))
        
        # 3通道tensor
        tensor_input = torch.randn(3, 64, 64)
        processed = processor.process(tensor_input)
        
        assert processed.shape == (1, 3, 32, 32)
    
    def test_process_grayscale_tensor(self):
        """测试处理灰度tensor"""
        processor = ImageProcessor(image_size=(32, 32))
        
        # 单通道tensor
        gray_tensor = torch.randn(64, 64)
        processed = processor.process(gray_tensor)
        
        assert processed.shape == (1, 3, 32, 32)  # 应该转换为3通道


class TestImageRecognizer:
    """测试图像识别器"""
    
    def test_recognizer_initialization(self):
        """测试识别器初始化"""
        recognizer = ImageRecognizer(num_classes=10)
        
        assert recognizer is not None
        assert hasattr(recognizer, 'model')
        assert hasattr(recognizer, 'processor')
    
    def test_recognize_numpy_array(self):
        """测试识别numpy数组"""
        recognizer = ImageRecognizer(num_classes=10)
        
        rgb_array = np.random.randint(0, 255, (32, 32, 3), dtype=np.uint8)
        result = recognizer.recognize(rgb_array)
        
        assert 'class' in result
        assert 'confidence' in result
        assert 'all_probabilities' in result
        assert isinstance(result['class'], int)
        assert 0 <= result['class'] < 10
        assert 0 <= result['confidence'] <= 1
    
    def test_recognize_tensor(self):
        """测试识别PyTorch tensor"""
        recognizer = ImageRecognizer(num_classes=10)
        
        tensor_input = torch.randn(3, 32, 32)
        result = recognizer.recognize(tensor_input)
        
        assert 'class' in result
        assert 'confidence' in result
        assert isinstance(result['class'], int)
    
    def test_recognize_batch(self):
        """测试批量识别"""
        recognizer = ImageRecognizer(num_classes=10)
        
        batch_inputs = [
            np.random.randint(0, 255, (32, 32, 3), dtype=np.uint8),
            torch.randn(3, 32, 32),
            np.random.randint(0, 255, (32, 32), dtype=np.uint8)
        ]
        
        results = recognizer.recognize(batch_inputs)
        
        assert isinstance(results, list)
        assert len(results) == 3
        
        for result in results:
            assert 'class' in result
            assert 'confidence' in result
    
    def test_model_in_eval_mode(self):
        """测试模型处于评估模式"""
        recognizer = ImageRecognizer(num_classes=10)
        
        assert not recognizer.model.training  # 应该是评估模式


if __name__ == "__main__":
    pytest.main([__file__, "-v"])